import numpy as np
from category_encoders import TargetEncoder
from imblearn.over_sampling import SMOTE
from sklearn.preprocessing import StandardScaler

def feature_engineering(X_train, X_test, y_train):
    """特征工程 + SMOTE + 标准化"""

    # 1️⃣ 目标编码
    categorical_features = ['BusinessTravel', 'Department', 'EducationField', 'JobRole', 'MaritalStatus', 'OverTime']
    encoder = TargetEncoder(cols=categorical_features)
    X_train[categorical_features] = encoder.fit_transform(X_train[categorical_features], y_train)
    X_test[categorical_features] = encoder.transform(X_test[categorical_features])

    # 2️⃣ 比例特征
    X_train['YearsAtCompany_Age_ratio'] = X_train['YearsAtCompany'] / (X_train['Age'] + 1e-5)
    X_test['YearsAtCompany_Age_ratio'] = X_test['YearsAtCompany'] / (X_test['Age'] + 1e-5)

    X_train['TotalWorkingYears_Age_ratio'] = X_train['TotalWorkingYears'] / (X_train['Age'] + 1e-5)
    X_test['TotalWorkingYears_Age_ratio'] = X_test['TotalWorkingYears'] / (X_test['Age'] + 1e-5)

    X_train['CurrentRole_Company_ratio'] = X_train['YearsInCurrentRole'] / (X_train['YearsAtCompany'] + 1e-5)
    X_test['CurrentRole_Company_ratio'] = X_test['YearsInCurrentRole'] / (X_test['YearsAtCompany'] + 1e-5)

    # 3️⃣ 对数变换
    for col in ['MonthlyIncome', 'DistanceFromHome', 'NumCompaniesWorked']:
        if col in X_train.columns:
            X_train[col] = np.log1p(X_train[col])
            X_test[col] = np.log1p(X_test[col])

    # 4️⃣ SMOTE 过采样
    sm = SMOTE(random_state=42)
    X_train_res, y_train_res = sm.fit_resample(X_train, y_train)

    # 5️⃣ 标准化
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train_res)
    X_test_scaled = scaler.transform(X_test)

    return X_train_scaled, X_test_scaled, y_train_res, encoder, scaler
